package com.ruoyi.ai.service.impl;

import cn.hutool.core.collection.CollectionUtil;
import cn.hutool.core.util.StrUtil;
import cn.hutool.json.JSONArray;
import cn.hutool.json.JSONUtil;
import com.ruoyi.ai.domain.Gpt35TurboVO;
import com.ruoyi.ai.domain.StreamParamVO;
import com.ruoyi.ai.service.IChatGtpService;
import com.ruoyi.chatgpt.domain.GptChat;
import com.ruoyi.chatgpt.domain.GptKeyManager;
import com.ruoyi.chatgpt.domain.GptMessage;
import com.ruoyi.chatgpt.service.IGptChatService;
import com.ruoyi.chatgpt.service.IGptKeyManagerService;
import com.ruoyi.chatgpt.service.IGptMessageService;
import com.ruoyi.common.constant.Constants;
import com.ruoyi.common.core.domain.entity.SysUser;
import com.ruoyi.common.core.redis.RedisLock;
import com.ruoyi.common.utils.SecurityUtils;
import com.ruoyi.common.utils.json.JsonUtil;
import com.ruoyi.system.service.ISysUserService;
import org.apache.commons.lang3.StringUtils;
import org.redisson.api.RLock;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Service;
import org.springframework.transaction.annotation.Transactional;

import javax.servlet.http.HttpServletResponse;
import java.io.*;
import java.net.HttpURLConnection;
import java.net.URL;
import java.util.*;
import java.util.concurrent.TimeUnit;

/**
 * @description:ChatGtpServiceImpl
 * @author: caoes
 * create:2023/6/30 15:42
 **/
@Service
public class ChatGtpServiceImpl implements IChatGtpService {

    @Autowired
    private RedisLock redisLock;
    @Autowired
    private ISysUserService iSysUserService;

    @Autowired(required = false)
    private RedisTemplate<Object, Object> redisTemplate;

    @Autowired
    private IGptChatService gptChatService;

    @Autowired
    private IGptMessageService gptMessageService;

    @Autowired
    private IGptKeyManagerService gptKeyManagerService;
    @Override
    @Transactional(rollbackFor = Exception.class)
    public void chatBotSseStream(StreamParamVO streamParamVO, HttpServletResponse response) {
        Long userId = SecurityUtils.getUserId();
        String lockName = "aks_" + userId;
        try {
            RLock lock = redisLock.getRLock(lockName);
            if (lock.isLocked()) {
                throw new RuntimeException("回复中");
            }
            redisLock.lock(lockName);
            response.setContentType("text/event-stream");
            response.setCharacterEncoding("utf-8");
            SysUser user = iSysUserService.selectUserById(userId);
            // 校验 key
            checkChatStream(user, streamParamVO);

            InputStream is = this.sendRequestBefore(user, streamParamVO);
            String line = null;
            String answerContent = "";
            BufferedReader reader = new BufferedReader(new InputStreamReader(is));
            PrintWriter pw = response.getWriter();
            while((line = reader.readLine()) != null) {
                //首先对行数据进行处理
                if (StrUtil.isNotBlank(line) ) {
                    line = CollectionUtil.removeEmpty(StrUtil.split(line, "data: ")).get(0);
                    if (!StrUtil.contains(line, "[DONE]")){
                        String oneWord = catchTextGpt(line);
                        if (StrUtil.isNotBlank(oneWord)){
                            answerContent = answerContent+oneWord;
                        }
                    }
                    pw.write(line);
                    pw.flush();
                    TimeUnit.MILLISECONDS.sleep(50);
                }
            }
            //处理完了后将次条聊天记录进行记录
            if(StrUtil.isNotBlank(answerContent)){
                //保存聊天记录
                saveAnswerContent(user, streamParamVO, answerContent);
            }
            is.close();
            pw.close();
            reader.close();
        } catch (Exception e) {
            throw new RuntimeException(e.getMessage());
        } finally {
            //解锁
            redisLock.unlock(lockName);
        }
    }
    private void saveAnswerContent(SysUser user, StreamParamVO streamParamVO, String answerContent) {
        Long chatId = streamParamVO.getChatId();
        Date now = new Date();

        if (Objects.isNull(streamParamVO.getChatId())) {
            // save chat
            GptChat chat = new GptChat();
            chat.setName(streamParamVO.getPrompt().length() > 20 ? streamParamVO.getPrompt().substring(0,20) : streamParamVO.getPrompt());
            chat.setStatus(1);
            chat.setUserId(user.getUserId());
            chat.setCreateTime(now);
            chat.setUpdateTime(now);
            gptChatService.insertGptChat(chat);

            chatId = chat.getId();
        }

        GptMessage gptMessage = new GptMessage();
        gptMessage.setChatId(chatId);
        gptMessage.setAskContent(streamParamVO.getPrompt());
        gptMessage.setAnswerContent(answerContent);
        gptMessage.setCreateTime(now);
        gptMessage.setUpdateTime(now);

        gptMessageService.insertGptMessage(gptMessage);
    }

    private InputStream sendRequestBefore(SysUser loginUser, StreamParamVO streamParamVO) throws Exception{
        InputStream in = null;
        String prompt = streamParamVO.getPrompt();
        if (StrUtil.isBlank(prompt)) {
            throw new RuntimeException("输入内容为空");
        }
        Object obj = redisTemplate.opsForValue().get(loginUser.getUserId()+"");
        if (!Objects.isNull(obj)) {
            throw new RuntimeException("正在回复");
        }
        String apikey = getOpenAiKey();
        if (StrUtil.isBlank(apikey)){
            throw new RuntimeException("无可用key");
        }

        redisTemplate.opsForValue().set(loginUser.getUserId()+"", true, 30, TimeUnit.SECONDS);
        List<Gpt35TurboVO> chatContext = getChatDigContext(loginUser.getUserId(), streamParamVO);
        Map<Object, Object> objectObjectHashMap = new HashMap<>();
        objectObjectHashMap.put("model", "gpt-3.5-turbo");
        objectObjectHashMap.put("messages", chatContext);
        objectObjectHashMap.put("stream", true);
        objectObjectHashMap.put("temperature", 0);
        objectObjectHashMap.put("frequency_penalty", 0);
        objectObjectHashMap.put("presence_penalty", 0.6);
        String bodyJson = JSONUtil.toJsonStr(objectObjectHashMap);
        URL url = new URL(Constants.OPEN_AI_REQUEST_URL);
        HttpURLConnection urlConnection = (HttpURLConnection)url.openConnection();
        urlConnection.setRequestMethod("POST");
        urlConnection.setDoOutput(true);
        urlConnection.setDoInput(true);
        urlConnection.setUseCaches(false);
        urlConnection.setRequestProperty("Connection", "Keep-Alive");
        urlConnection.setRequestProperty("Charset", "UTF-8");
        urlConnection.setRequestProperty("Authorization", "Bearer " + apikey);
        urlConnection.setRequestProperty("Content-Type", "application/json; charset=UTF-8");
        byte[] dataBytes = bodyJson.getBytes();
        urlConnection.setRequestProperty("Content-Length", String.valueOf(dataBytes.length));
        OutputStream os = urlConnection.getOutputStream();
        os.write(dataBytes);
        in =   new BufferedInputStream(urlConnection.getInputStream());
        os.flush();
        os.close();
        return in;
    }

    private String getOpenAiKey() {
        GptKeyManager keyManager = new GptKeyManager();
        keyManager.setStatus(1);
        List<GptKeyManager> keyManagerList = gptKeyManagerService.selectGptKeyManagerList(keyManager);
        //判断是否key额度用完
        if (CollectionUtil.isEmpty(keyManagerList) || keyManagerList.size() <= 0) {
            throw new RuntimeException("key额度耗尽");
        }
        //获取第一个key,然后将第一个key存入缓存
        String apikey = keyManagerList.get(0).getSecretKey();
        redisTemplate.opsForValue().set("apikey", apikey);
        //检查key
        checkKey(apikey);
        return apikey;
    }

    private void checkKey(String apikey) {
        // 检查key 有效

    }

    private List<Gpt35TurboVO> getChatDigContext(Long userId, StreamParamVO streamParamVO) {
        List<Gpt35TurboVO> messages = new ArrayList<>();
        //首先获取角色,默认角色
        String default_role = "我是一个全能机器人,可以回答你任何问题";
        //设定系统所扮演的角色
        Gpt35TurboVO gpt35TurboVOSys = new Gpt35TurboVO();
        gpt35TurboVOSys.setRole("system");
        gpt35TurboVOSys.setContent(default_role);
        messages.add(gpt35TurboVOSys);
        // chatId 是否为空
        if(!Objects.isNull(streamParamVO.getChatId())) {
            GptChat gptChat = gptChatService.selectGptChatById(streamParamVO.getChatId());
            if (gptChat == null || !gptChat.getUserId().equals(userId)) {
                throw new RuntimeException("对话无效");
            }
            List<GptMessage> gptMessageList = gptMessageService.queryGptMessageList(streamParamVO.getChatId());
            if (CollectionUtil.isNotEmpty(gptMessageList)) {
                for (GptMessage message : gptMessageList) {
                    Gpt35TurboVO gpt35TurboUser = new Gpt35TurboVO();
                    //用户询问的问题
                    gpt35TurboUser.setRole("user");
                    gpt35TurboUser.setContent(message.getAskContent());
                    messages.add(gpt35TurboUser);
                    //机器人回答的问题
                    Gpt35TurboVO gpt35TurAssistant = new Gpt35TurboVO();
                    gpt35TurAssistant.setRole("assistant");
                    gpt35TurAssistant.setContent(message.getAnswerContent());
                    messages.add(gpt35TurAssistant);
                }
            }
        }

        //最后查询用户询问的问题
        Gpt35TurboVO gpt35TurboUser = new Gpt35TurboVO();
        gpt35TurboUser.setRole("user");
        gpt35TurboUser.setContent(streamParamVO.getPrompt());
        messages.add(gpt35TurboUser);
        return messages;
    }

    /**
     * 处理单独打印的文字
     * @param str
     * @return
     */
    public String catchTextGpt(String str){
        String choices = JsonUtil.parseMiddleData(str, "choices");
        JSONArray jsonArray = JSONUtil.parseArray(choices);
        String string = jsonArray.getByPath("[0].delta").toString();
        String content = JsonUtil.parseMiddleData(string, "content");
        return content;
    }

    private void checkChatStream(SysUser user, StreamParamVO streamParamVO) {

        if (StringUtils.isBlank(streamParamVO.getPrompt())) {
            throw new RuntimeException("内容不可为空");
        }
        //判断是否存在违禁词
        this.catchDisableWordRegx(streamParamVO.getPrompt());
        // 用户字数判断
        if (user == null || user.getQuantity() - streamParamVO.getPrompt().length() < 0) {
            throw new RuntimeException("额度不足,请先充值后在使用");
        }
    }

    private void catchDisableWordRegx(String prompt) {
    }
}
